import torch
from torch.utils.data import Dataset
from transformers import Trainer, TrainingArguments


# 自定义 PyTorch 模型
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = torch.nn.Linear(768, 2)  # 示例中输入维度为768，输出维度为2

    def forward(self, input_ids):
        return self.fc(input_ids)


# 自定义损失函数
def custom_loss(predictions, labels):
    # 在这里定义你的损失计算逻辑
    return torch.nn.functional.cross_entropy(predictions, labels)


# 自定义数据集
class CustomDataset(Dataset):
    def __init__(self, inputs, labels):
        self.inputs = inputs
        self.labels = labels

    def __len__(self):
        return len(self.inputs)

    def __getitem__(self, idx):
        return {'input_ids': self.inputs[idx], 'labels': self.labels[idx]}


# 模拟数据
train_inputs = torch.randn(100, 768)  # 示例中的输入维度为768
train_labels = torch.randint(2, (100,))  # 示例中的标签为0或1
train_dataset = CustomDataset(train_inputs, train_labels)

# 定义训练参数
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=8,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    compute_loss=custom_loss,  # 指定自定义损失函数
)

# 初始化模型和 Trainer
model = SimpleModel()
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

# 开始训练
trainer.train()
